{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Traffic light state learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Jupyter notebook for training a tensorflow cnn model capable to detect traffic light color. The training is based on cropped images from Udacity's simulator, the images provided in the rosbag, the Bosch dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from scipy import misc\n",
    "from random import shuffle, random\n",
    "import cv2\n",
    "import os\n",
    "import six.moves.urllib as urllib\n",
    "import tarfile\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "from time import gmtime, strftime\n",
    "import yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_training_examples(directory='small_lights', sub_directory='.', light=0):\n",
    "    \"\"\"Extracting filenames and predefined labels from a directory\"\"\"\n",
    "    labels, filenames = (list(), list())\n",
    "    path = os.path.join(os.getcwd(), directory, sub_directory)\n",
    "    examples = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]\n",
    "    labels += [light] * len(examples)\n",
    "    filenames += examples\n",
    "    return np.array(filenames), np.array(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def read_image_from_disk(filename, flipping=True):\n",
    "    \"\"\"\n",
    "    Reads a png image from disk and \n",
    "    converts it into a Numpy ndarray\n",
    "    \"\"\"\n",
    "    file_contents = misc.imread(filename, mode='RGB') # The model expects RGB images\n",
    "    if file_contents.shape != (32, 64):\n",
    "        file_contents = misc.imresize(file_contents, (32, 64)) # The model expects 32x64x3 images\n",
    "    if random() > 0.5 and flipping is True:\n",
    "        return np.fliplr(file_contents)\n",
    "    else: \n",
    "        return file_contents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def chunks(l, n):\n",
    "    \"\"\"Yield successive n-sized chunks from l.\"\"\"\n",
    "    for i in range(0, len(l), n):\n",
    "        yield l[i:i + n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def shuffle_data(filenames, labels):\n",
    "    \"\"\"Shuffles Numpy ndarrays for filenames and labels\"\"\"\n",
    "    index = list(range(len(filenames)))\n",
    "    shuffle(index)\n",
    "    return filenames[index], labels[index]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def batch_feed(filenames, labels, batch_size=5, shuffling=True):\n",
    "    \"\"\"Batch procedure for feeding the training process\"\"\"\n",
    "    index = list(range(len(filenames)))\n",
    "    if shuffling:\n",
    "        shuffle(index)\n",
    "    for batch in chunks(index, batch_size):\n",
    "        batch_labels = labels[batch]\n",
    "        batch_train = np.array([read_image_from_disk(file) for file in filenames[batch]]) / 255.\n",
    "        yield(batch_train, batch_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def create_placeholders():\n",
    "    \"\"\"Creates i/O and dropout placeholders\"\"\"\n",
    "    keep_prob = tf.placeholder(tf.float32, name=\"keep_prob\")\n",
    "    input_layer = tf.placeholder(tf.float32, shape=(None, None, None, Channels), name=\"input_layer\")\n",
    "    output_layer = tf.placeholder(tf.int32, shape=(None))\n",
    "    return input_layer, keep_prob, output_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def create_layers(input_layer, keep_prob, height, width, classes):\n",
    "    \"\"\"    \n",
    "    ____________________________________________________________________________________________________\n",
    "    Layer (type)                     Output Shape          Param #     Connected to\n",
    "    ====================================================================================================\n",
    "    convolution2d_1 (Convolution2D)  (None, 112, 112, 16)  448         convolution2d_input_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    maxpooling2d_1 (MaxPooling2D)    (None, 37, 37, 16)    0           convolution2d_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dropout_1 (Dropout)              (None, 37, 37, 16)    0           maxpooling2d_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    convolution2d_2 (Convolution2D)  (None, 37, 37, 32)    4640        dropout_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    maxpooling2d_2 (MaxPooling2D)    (None, 12, 12, 32)    0           convolution2d_2[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dropout_2 (Dropout)              (None, 12, 12, 32)    0           maxpooling2d_2[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    convolution2d_3 (Convolution2D)  (None, 12, 12, 64)    18496       dropout_2[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    maxpooling2d_3 (MaxPooling2D)    (None, 6, 6, 64)      0           convolution2d_3[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dropout_3 (Dropout)              (None, 6, 6, 64)      0           maxpooling2d_3[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    flatten_1 (Flatten)              (None, 2304)          0           dropout_3[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dense_1 (Dense)                  (None, 128)           295040      flatten_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dropout_4 (Dropout)              (None, 128)           0           dense_1[0][0]\n",
    "    ____________________________________________________________________________________________________\n",
    "    dense_2 (Dense)                  (None, 3)             387         dropout_4[0][0]\n",
    "    ====================================================================================================\n",
    "    Total params: 319,011\n",
    "    Trainable params: 319,011\n",
    "    Non-trainable params: 0\n",
    "    \n",
    "    \"\"\"\n",
    "    resized_input = tf.image.resize_images(images = input_layer, \n",
    "                                           size = (height, width)\n",
    "                                          )\n",
    "    \n",
    "    conv1 = tf.layers.conv2d(inputs=resized_input,\n",
    "                             filters=16, kernel_size=[3, 3],\n",
    "                             padding=\"same\", activation=tf.nn.relu)\n",
    "\n",
    "    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[3, 3],\n",
    "                                    strides=2)\n",
    "\n",
    "    dropout1 = tf.nn.dropout(pool1, keep_prob)\n",
    "    \n",
    "    conv2 = tf.layers.conv2d(inputs=dropout1, \n",
    "                             filters=32, kernel_size=[3, 3], \n",
    "                             padding=\"same\", activation=tf.nn.relu)\n",
    "\n",
    "    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[3, 3],\n",
    "                                    strides=2)\n",
    "    \n",
    "    dropout2 = tf.nn.dropout(pool2, keep_prob)\n",
    "    \n",
    "    conv3 = tf.layers.conv2d(inputs=dropout2, \n",
    "                             filters=64, kernel_size=[3, 3], \n",
    "                             padding=\"same\", activation=tf.nn.relu)\n",
    "\n",
    "    pool3 = tf.layers.max_pooling2d(inputs=conv3, pool_size=[2, 2],\n",
    "                                    strides=2)\n",
    "    \n",
    "    dropout3 = tf.nn.dropout(pool3, keep_prob)\n",
    "     \n",
    "    flat = tf.contrib.layers.flatten(dropout3)\n",
    "    \n",
    "    connected = tf.layers.dense(inputs=flat, units=128,\n",
    "                                  activation=tf.tanh)\n",
    "    \n",
    "    dropout4 = tf.nn.dropout(connected, keep_prob-0.1)\n",
    "    \n",
    "    logits = tf.layers.dense(inputs=dropout4, units=classes,\n",
    "                             name=\"output_logits\")\n",
    "    \n",
    "    return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def optimization(logits):\n",
    "    \"\"\"Optimization based on AdamOptimizer\"\"\"\n",
    "    one_hot_y = tf.one_hot(output_layer, Classes)\n",
    "    loss_operation = tf.reduce_mean(\n",
    "        tf.nn.softmax_cross_entropy_with_logits(logits=logits,\n",
    "                                                labels=one_hot_y)\n",
    "    )\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=0.0003)\n",
    "    training_operation = optimizer.minimize(loss_operation)\n",
    "    correct_prediction = tf.equal(\n",
    "        tf.argmax(logits, 1, name=\"output_class\"),\n",
    "        tf.argmax(one_hot_y, 1))\n",
    "    accuracy_operation = tf.reduce_mean(\n",
    "        tf.cast(correct_prediction, tf.float32))\n",
    "    return loss_operation, training_operation, accuracy_operation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Collecting available examples from separated directories\n",
    "sim_green_files,     sim_green_labels     = get_training_examples(sub_directory='green_sim',     light=2)\n",
    "sim_red_files,       sim_red_labels       = get_training_examples(sub_directory='red_sim',       light=0)\n",
    "sim_yellow_files,    sim_yellow_labels    = get_training_examples(sub_directory='yellow_sim',    light=1)\n",
    "sim_none_files,      sim_none_labels      = get_training_examples(sub_directory='none_sim',      light=3)\n",
    "bosch_green_files,   bosch_green_labels   = get_training_examples(sub_directory='green_bosch',   light=2)\n",
    "bosch_red_files,     bosch_red_labels     = get_training_examples(sub_directory='red_bosch',     light=0)\n",
    "bosch_yellow_files,  bosch_yellow_labels  = get_training_examples(sub_directory='yellow_bosch',  light=1)\n",
    "bosch_none_files,    bosch_none_labels    = get_training_examples(sub_directory='none_bosch',    light=3)\n",
    "rosbag_green_files,  rosbag_green_labels  = get_training_examples(sub_directory='green_rosbag',  light=2)\n",
    "rosbag_red_files,    rosbag_red_labels    = get_training_examples(sub_directory='red_rosbag',    light=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Putting together all the examples for the training\n",
    "filenames = np.concatenate((sim_green_files, sim_red_files, sim_yellow_files, sim_none_files,\n",
    "                        bosch_green_files, bosch_red_files, bosch_yellow_files, bosch_none_files,\n",
    "                        rosbag_green_files, rosbag_red_files))\n",
    "labels = np.concatenate((sim_green_labels, sim_red_labels, sim_yellow_labels, sim_none_labels,\n",
    "                        bosch_green_labels, bosch_red_labels, bosch_yellow_labels, bosch_none_labels,\n",
    "                        rosbag_green_labels, rosbag_red_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0, 1, 2, 3]), array([4212,  312, 5720, 5358], dtype=int64))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Reporting label distributions\n",
    "np.unique(labels, return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train examples: 14041\n",
      "Validation examples 1561\n",
      "Epoch: 0 Batch: 0 Train accuracy 0.400\n",
      "Epoch: 0 Batch: 2808 Train accuracy 0.945 Validation accuracy 0.961\n",
      "Model saved\n",
      "Epoch: 1 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 1 Batch: 2808 Train accuracy 0.974 Validation accuracy 0.978\n",
      "Model saved\n",
      "Epoch: 2 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 2 Batch: 2808 Train accuracy 0.979 Validation accuracy 0.980\n",
      "Model saved\n",
      "Epoch: 3 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 3 Batch: 2808 Train accuracy 0.980 Validation accuracy 0.983\n",
      "Model saved\n",
      "Epoch: 4 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 4 Batch: 2808 Train accuracy 0.981 Validation accuracy 0.983\n",
      "Epoch: 5 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 5 Batch: 2808 Train accuracy 0.981 Validation accuracy 0.982\n",
      "Epoch: 6 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 6 Batch: 2808 Train accuracy 0.982 Validation accuracy 0.980\n",
      "Epoch: 7 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 7 Batch: 2808 Train accuracy 0.984 Validation accuracy 0.984\n",
      "Model saved\n",
      "Epoch: 8 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 8 Batch: 2808 Train accuracy 0.984 Validation accuracy 0.983\n",
      "Epoch: 9 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 9 Batch: 2808 Train accuracy 0.984 Validation accuracy 0.983\n",
      "Epoch: 10 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 10 Batch: 2808 Train accuracy 0.985 Validation accuracy 0.986\n",
      "Model saved\n",
      "Epoch: 11 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 11 Batch: 2808 Train accuracy 0.985 Validation accuracy 0.986\n",
      "Model saved\n",
      "Epoch: 12 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 12 Batch: 2808 Train accuracy 0.987 Validation accuracy 0.985\n",
      "Epoch: 13 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 13 Batch: 2808 Train accuracy 0.985 Validation accuracy 0.982\n",
      "Epoch: 14 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 14 Batch: 2808 Train accuracy 0.985 Validation accuracy 0.985\n",
      "Epoch: 15 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 15 Batch: 2808 Train accuracy 0.986 Validation accuracy 0.987\n",
      "Model saved\n",
      "Epoch: 16 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 16 Batch: 2808 Train accuracy 0.986 Validation accuracy 0.984\n",
      "Epoch: 17 Batch: 0 Train accuracy 0.800\n",
      "Epoch: 17 Batch: 2808 Train accuracy 0.987 Validation accuracy 0.980\n",
      "Epoch: 18 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 18 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.981\n",
      "Epoch: 19 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 19 Batch: 2808 Train accuracy 0.987 Validation accuracy 0.983\n",
      "Epoch: 20 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 20 Batch: 2808 Train accuracy 0.987 Validation accuracy 0.987\n",
      "Model saved\n",
      "Epoch: 21 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 21 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.986\n",
      "Epoch: 22 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 22 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.985\n",
      "Epoch: 23 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 23 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.988\n",
      "Model saved\n",
      "Epoch: 24 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 24 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.982\n",
      "Epoch: 25 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 25 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.985\n",
      "Epoch: 26 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 26 Batch: 2808 Train accuracy 0.987 Validation accuracy 0.985\n",
      "Epoch: 27 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 27 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.981\n",
      "Epoch: 28 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 28 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.986\n",
      "Epoch: 29 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 29 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.986\n",
      "Epoch: 30 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 30 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.986\n",
      "Epoch: 31 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 31 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.987\n",
      "Epoch: 32 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 32 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.984\n",
      "Epoch: 33 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 33 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.989\n",
      "Model saved\n",
      "Epoch: 34 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 34 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.986\n",
      "Epoch: 35 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 35 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.986\n",
      "Epoch: 36 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 36 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.986\n",
      "Epoch: 37 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 37 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.985\n",
      "Epoch: 38 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 38 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.985\n",
      "Epoch: 39 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 39 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.985\n",
      "Epoch: 40 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 40 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.986\n",
      "Epoch: 41 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 41 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.983\n",
      "Epoch: 42 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 42 Batch: 2808 Train accuracy 0.988 Validation accuracy 0.985\n",
      "Epoch: 43 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 43 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.987\n",
      "Epoch: 44 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 44 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.988\n",
      "Epoch: 45 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 45 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.988\n",
      "Epoch: 46 Batch: 0 Train accuracy 0.800\n",
      "Epoch: 46 Batch: 2808 Train accuracy 0.991 Validation accuracy 0.984\n",
      "Epoch: 47 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 47 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.986\n",
      "Epoch: 48 Batch: 0 Train accuracy 0.800\n",
      "Epoch: 48 Batch: 2808 Train accuracy 0.989 Validation accuracy 0.984\n",
      "Epoch: 49 Batch: 0 Train accuracy 1.000\n",
      "Epoch: 49 Batch: 2808 Train accuracy 0.990 Validation accuracy 0.984\n",
      "best epoch at ./state_detection-33\n"
     ]
    }
   ],
   "source": [
    "Height, Width, Channels = (32, 64, 3)\n",
    "Classes = 4\n",
    "\n",
    "epochs = 50\n",
    "batch_size = 5\n",
    "validation = 0.10\n",
    "best_result = 0.0\n",
    "\n",
    "tf.reset_default_graph() \n",
    "\n",
    "input_layer, keep_prob, output_layer = create_placeholders()\n",
    "logits = create_layers(input_layer, keep_prob, Height, Width, Classes)\n",
    "loss_operation, training_operation, accuracy_operation = optimization(logits)\n",
    "saver = tf.train.Saver(max_to_keep=10)\n",
    "\n",
    "filenames, labels = shuffle_data(filenames, labels)\n",
    "\n",
    "number_examples = len(labels)\n",
    "train_space = int(number_examples * (1.0 - validation))\n",
    "train_filenames, train_labels = filenames[:train_space], labels[:train_space]\n",
    "validation_filenames, validation_labels = filenames[train_space:], labels[train_space:]\n",
    "\n",
    "print(\"Train examples: %i\" % (len(train_labels)))\n",
    "print(\"Validation examples %i\" % (len(validation_labels)))\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    for e in range(epochs):\n",
    "        train_loss, train_accuracy = 0.0, 0.0\n",
    "        for n, (X, y) in enumerate(batch_feed(train_filenames, train_labels, batch_size=batch_size, shuffling=True)):\n",
    "            _, loss, accuracy = sess.run([training_operation, loss_operation, accuracy_operation],\n",
    "                                      feed_dict={input_layer: X,\n",
    "                                                 output_layer: y,\n",
    "                                                 keep_prob: 0.8\n",
    "                                                })\n",
    "            train_loss += loss\n",
    "            train_accuracy += accuracy\n",
    "            if n % 5000 == 0:\n",
    "                print(\"Epoch: %i Batch: %i Train accuracy %0.3f\" % (e, n, train_accuracy / (n + 1.0)))\n",
    "        \n",
    "        train_loss = train_loss / (n + 1.0)\n",
    "        train_accuracy = train_accuracy / (n + 1.0)\n",
    "            \n",
    "        validation_loss, validation_accuracy = 0.0, 0.0\n",
    "        for m, (X, y) in enumerate(batch_feed(validation_filenames, validation_labels, batch_size=100, shuffling=False)):\n",
    "            loss, accuracy = sess.run([loss_operation, accuracy_operation],\n",
    "                                      feed_dict={input_layer: X,\n",
    "                                                 output_layer: y,\n",
    "                                                 keep_prob: 1.0})\n",
    "            validation_loss += loss\n",
    "            validation_accuracy += accuracy\n",
    "        \n",
    "        validation_loss = validation_loss / (m + 1.0)\n",
    "        validation_accuracy = validation_accuracy / (m + 1.0)\n",
    "\n",
    "        print(\"Epoch: %i Batch: %i Train accuracy %0.3f Validation accuracy %0.3f\" % (e, n, train_accuracy, validation_accuracy))\n",
    "        \n",
    "        # Early cutoff - Keeping only the epoch with best validation accuracy\n",
    "        if validation_accuracy > best_result:\n",
    "            best_result = validation_accuracy\n",
    "            path = saver.save(sess, './state_detection', global_step=e)\n",
    "            print(\"Model saved\")\n",
    "            best_epoch_path = path\n",
    "            \n",
    "print(\"best epoch at {}\".format(best_epoch_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Traffic light state testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state_detection-33\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print (best_epoch_path)\n",
    "except:\n",
    "    best_epoch_path = 'state_detection-33'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from state_detection-33\n"
     ]
    }
   ],
   "source": [
    "# Recovering the neural network\n",
    "session = tf.Session()\n",
    "loader = tf.train.import_meta_graph(best_epoch_path+'.meta')\n",
    "loader.restore(session, best_epoch_path)\n",
    "graph = tf.get_default_graph()\n",
    "out = graph.get_tensor_by_name(\"output_class:0\")\n",
    "input_layer = graph.get_tensor_by_name(\"input_layer:0\")\n",
    "keep_prob = graph.get_tensor_by_name(\"keep_prob:0\")\n",
    "counter = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_training_examples(directory='small_lights', sub_directory='.', light=0):\n",
    "    \"\"\"Extracting filenames and predefined labels from a directory\"\"\"\n",
    "    labels, filenames = (list(), list())\n",
    "    path = os.path.join(os.getcwd(), directory, sub_directory)\n",
    "    examples = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]\n",
    "    labels += [light] * len(examples)\n",
    "    filenames += examples\n",
    "    return np.array(filenames), np.array(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def read_image_from_disk(filename, flipping=True):\n",
    "    \"\"\"\n",
    "    Reads a png image from disk and \n",
    "    converts it into a Numpy ndarray\n",
    "    \"\"\"\n",
    "    file_contents = misc.imread(filename, mode='RGB')\n",
    "    if file_contents.shape != (32, 64):\n",
    "        file_contents = misc.imresize(file_contents, (32, 64))\n",
    "    if random() > 0.5 and flipping is True:\n",
    "        return np.fliplr(file_contents)\n",
    "    else: \n",
    "        return file_contents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collecting available examples from separated directories\n",
    "sim_green_files,     sim_green_labels     = get_training_examples(sub_directory='green_sim',     light=2)\n",
    "sim_red_files,       sim_red_labels       = get_training_examples(sub_directory='red_sim',       light=0)\n",
    "sim_yellow_files,    sim_yellow_labels    = get_training_examples(sub_directory='yellow_sim',    light=1)\n",
    "sim_none_files,      sim_none_labels      = get_training_examples(sub_directory='none_sim',      light=3)\n",
    "bosch_green_files,   bosch_green_labels   = get_training_examples(sub_directory='green_bosch',   light=2)\n",
    "bosch_red_files,     bosch_red_labels     = get_training_examples(sub_directory='red_bosch',     light=0)\n",
    "bosch_yellow_files,  bosch_yellow_labels  = get_training_examples(sub_directory='yellow_bosch',  light=1)\n",
    "bosch_none_files,    bosch_none_labels    = get_training_examples(sub_directory='none_bosch',    light=3)\n",
    "rosbag_green_files,  rosbag_green_labels  = get_training_examples(sub_directory='green_rosbag',  light=2)\n",
    "rosbag_red_files,    rosbag_red_labels    = get_training_examples(sub_directory='red_rosbag',    light=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# getting all the available images\n",
    "files = np.concatenate((sim_green_files, sim_red_files, sim_yellow_files, sim_none_files,\n",
    "                        bosch_green_files, bosch_red_files, bosch_yellow_files, bosch_none_files,\n",
    "                        rosbag_green_files, rosbag_red_files))\n",
    "labels = np.concatenate((sim_green_labels, sim_red_labels, sim_yellow_labels, sim_none_labels,\n",
    "                        bosch_green_labels, bosch_red_labels, bosch_yellow_labels, bosch_none_labels,\n",
    "                        rosbag_green_labels, rosbag_red_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "correct = 0\n",
    "wrong = 0\n",
    "for file, label in zip(files, labels):\n",
    "    img = read_image_from_disk(file, flipping=False)\n",
    "    result = session.run([out], feed_dict={input_layer: [img /255.], keep_prob: 1.0})\n",
    "    if (label ==result[0][0]):\n",
    "        correct += 1\n",
    "    else:\n",
    "        wrong += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.991\n"
     ]
    }
   ],
   "source": [
    "# This estimation contains training examples\n",
    "print (\"Accuracy: %0.3f\" % (float(correct)/(correct+wrong)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# getting all the simulator and rosbag images\n",
    "files = np.concatenate((sim_green_files, sim_red_files, sim_yellow_files, sim_none_files,\n",
    "                        rosbag_green_files, rosbag_red_files))\n",
    "labels = np.concatenate((sim_green_labels, sim_red_labels, sim_yellow_labels, sim_none_labels,\n",
    "                        rosbag_green_labels, rosbag_red_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "correct = 0\n",
    "wrong = 0\n",
    "for file, label in zip(files, labels):\n",
    "    img = read_image_from_disk(file, flipping=False)\n",
    "    result = session.run([out], feed_dict={input_layer: [img /255.], keep_prob: 1.0})\n",
    "    if (label ==result[0][0]):\n",
    "        correct += 1\n",
    "    else:\n",
    "        wrong += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.995\n"
     ]
    }
   ],
   "source": [
    "# This estimation contains training examples\n",
    "print (\"Accuracy: %0.3f\" % (float(correct)/(correct+wrong)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Labels in order to decode the answer\n",
    "labels = {0:'RED', 1:'YELLOW', 2:'GREEN', 3:'NONE'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GREEN\n"
     ]
    }
   ],
   "source": [
    "# Testing a couple of difficult examples\n",
    "img = read_image_from_disk(\"cnn_failing_RED1.jpg\", flipping=False)\n",
    "result = session.run([out], feed_dict={input_layer: [img /255.], keep_prob: 1.0})\n",
    "print (labels[result[0][0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GREEN\n"
     ]
    }
   ],
   "source": [
    "img = read_image_from_disk(\"cnn_failing_RED0.jpg\", flipping=False)\n",
    "result = session.run([out], feed_dict={input_layer: [img /255.], keep_prob: 1.0})\n",
    "print (labels[result[0][0]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
